/*
 * Wazuh Vulnerability scanner - Scan Orchestrator
 * Copyright (C) 2015, Wazuh Inc.
 * March 11, 2024.
 *
 * This program is free software; you can redistribute it
 * and/or modify it under the terms of the GNU General Public
 * License (version 2) as published by the FSF - Free Software
 * Foundation.
 */

#ifndef _BUILD_SINGLE_AGENT_INFO_CONTEXT_HPP
#define _BUILD_SINGLE_AGENT_INFO_CONTEXT_HPP

#include "chainOfResponsability.hpp"
#include "loggerHelper.h"
#include "scanContext.hpp"
#include "socketDBWrapper.hpp"
#include "vulnerabilityScanner.hpp"
#include "wazuhDBQueryBuilder.hpp"

/**
 * @brief Orchestrates queries over the global Wazuh system
 *
 * This class is responsible for managing the execution of queries within the global Wazuh environment.
 *
 * @tparam TScanContext scan context type.
 * @tparam TSocketDBWrapper socket database wrapper type
 */
template<typename TScanContext = ScanContext, typename TSocketDBWrapper = SocketDBWrapper>
class TBuildSingleAgentListInfoContext final : public AbstractHandler<std::shared_ptr<TScanContext>>
{

public:
    /**
     * @brief Construct a new global fetch object
     */
    explicit TBuildSingleAgentListInfoContext() = default;

    /**
     * @brief Handles request and passes control to the next step of the chain.
     *
     * @param data A shared pointer to the input data containing details of the query request.
     *
     * @return A shared pointer to a `TScanContext` object representing the query response.
     */
    std::shared_ptr<TScanContext> handleRequest(std::shared_ptr<TScanContext> data) override
    {
        nlohmann::json response;

        try
        {
            // Execute query
            TSocketDBWrapper::instance().query(
                WazuhDBQueryBuilder::builder()
                    .globalGetCommand(std::string("agent-info ") + data->agentId().data())
                    .build(),
                response);
        }
        catch (const SocketDbWrapperException& e)
        {
            throw WdbDataException(e.what(), data->agentId().data());
        }
        catch (std::exception& e)
        {
            throw std::runtime_error(std::string("Unable to retrieve agent-info (agent ") + data->agentId().data() +
                                     "). Reason: " + e.what() + ".");
        }

        // Return elements should be one agent.
        if (response.size() == 1)
        {
            const auto clusterStatus = data->clusterStatus();
            const auto clusterNodeName = data->clusterNodeName();
            const auto agentId = data->agentId();

            try
            {
                const auto& agent = response.front();
                const auto& wdbAgentNodeName = agent.at("node_name").get_ref<const std::string&>();
                const auto& wdbAgentConnectionStatus = agent.at("connection_status").get_ref<const std::string&>();
                const auto& wdbAgentVersion = agent.at("version").get_ref<const std::string&>();
                const auto& wdbAgentName = agent.at("name").get_ref<const std::string&>();
                const auto& wdbAgentIp = agent.at("ip").get_ref<const std::string&>();

                // Add the agent to the scan list if one of the conditions is met:
                //   - It's not in a cluster
                //   - It's in the cluster and matches the server node name and the connection status is active.
                if (!clusterStatus || (clusterNodeName == wdbAgentNodeName && "active" == wdbAgentConnectionStatus))
                {
                    data->m_agents.push_back(
                        {agentId.data(), wdbAgentName, Utils::leftTrim(wdbAgentVersion, "Wazuh "), wdbAgentIp});
                    logDebug2(
                        WM_VULNSCAN_LOGTAG, "Agent %s added to the list of agents to be re-scanned", agentId.data());
                }
                else
                {
                    logDebug2(WM_VULNSCAN_LOGTAG,
                              "Agent %s not added to the list of agents to be re-scanned (cluster status: %s) (cluster "
                              "node_name: %s) (connection_status: %s) (agent node_name: %s)",
                              agentId.data(),
                              clusterStatus ? "true" : "false",
                              clusterNodeName.data(),
                              wdbAgentConnectionStatus.c_str(),
                              wdbAgentNodeName.c_str());
                }
            }
            catch (const std::exception& e)
            {
                throw std::runtime_error("Invalid agent-info response format");
            }
        }
        return AbstractHandler<std::shared_ptr<TScanContext>>::handleRequest(std::move(data));
    }
};

using BuildSingleAgentListInfoContext = TBuildSingleAgentListInfoContext<>;

#endif // _BUILD_SINGLE_AGENT_INFO_CONTEXT_HPP
